package com.tlzh.controller;


import com.tlzh.pojo.dto.ChatRequest;
import jakarta.annotation.Resource;
import org.springframework.http.MediaType;
import org.springframework.web.bind.annotation.*;
import org.springframework.web.reactive.function.client.WebClient;
import org.springframework.web.servlet.mvc.method.annotation.SseEmitter;
import reactor.core.publisher.Flux;
import reactor.core.publisher.Mono;

import java.io.IOException;
import java.time.LocalDateTime;
import java.util.concurrent.atomic.AtomicBoolean;

/**
*  @新增加可打断controller
 *  @单线程打断
* */
@RestController
@RequestMapping("/chat")
@CrossOrigin
public class ChatV1Controller {
    private final AtomicBoolean isInterrupted = new AtomicBoolean(false);


    @Resource
    private WebClient webClient;

    @PostMapping(value = "/ai/ollamaChatStream", produces = MediaType.TEXT_EVENT_STREAM_VALUE)
    public SseEmitter chatModel(@RequestBody ChatRequest request) {
        SseEmitter emitter = new SseEmitter();
        isInterrupted.set(false);

        // 实际的 API 地址，你需要根据实际情况修改
        String apiUrl = "http://192.168.0.188:11434/api/chat";

        // logger 获取 请求时间
        System.out.println("Request received at: " + LocalDateTime.now());
        // 使用 WebClient 发送异步 POST 请求并获取 Flux<String>
        Flux<String> responseFlux = webClient.post()
                .uri(apiUrl)
                .contentType(MediaType.APPLICATION_JSON)
                .bodyValue(request)
                .retrieve()
                .bodyToFlux(String.class);
        System.out.println("Request Endt: " + LocalDateTime.now());


        System.out.println("Mono<Void> mono: " + LocalDateTime.now());
        // 订阅 Flux 并将数据发送到 SseEmitter
        Mono<Void> mono = responseFlux
                .takeWhile(part -> !isInterrupted.get()) // 当 isInterrupted 为 true 时停止发送数据
                .doOnNext(part -> {
                    try {
                        SseEmitter.SseEventBuilder event = SseEmitter.event()
                                .data(part)
                                .id(String.valueOf(System.currentTimeMillis()))
                                .reconnectTime(5000);
                        System.out.println("event Start: " + LocalDateTime.now());
                        emitter.send(event);
                        System.out.println("event End: " + LocalDateTime.now());
                    } catch (IOException e) {
                        emitter.completeWithError(e);
                    }
                })
                .doOnComplete(emitter::complete)
                .doOnError(emitter::completeWithError)
                .then();

        // 启动异步处理
        mono.subscribe();

        return emitter;
    }

    @PostMapping("/ai/interruptOllamaChatStream")
    public void interruptStreaming() {
        isInterrupted.set(true);
    }


}
